%% OPTIMAL POLICY SOLVER FOR OPTIMAL CONTRACT
function  [C, Beta, Eta, K, f, ExRet, mu_v, mu_y, vol_y, vol_phi] = optimal_optimal_policy(V0, C0, Beta0, Eta0, params, grid_params)

%% parameters
temp = num2cell(params);
[rho, r, ra, alpha, sigma, lambda, M, eta_min, eta_max] = temp{:};

[n_grid, Y_mat, PHI_mat, dy, dphi] = grid_params{:};

%% Value function and derivatives
V1 = [V0(1,:);V0;V0(end,:)];
V = [V1(:,1),V1, V1(:,end)];

%Initialize matrices
Vy_f = zeros(n_grid+2,n_grid+2);
Vy_b = zeros(n_grid+2,n_grid+2);
Vy_c = zeros(n_grid+2, n_grid+2);

Vphi_f = zeros(n_grid+2,n_grid+2);
Vphi_b = zeros(n_grid+2,n_grid+2);
Vphi_c = zeros(n_grid+2, n_grid+2);

Vyphi = zeros(n_grid+2,n_grid+2);

% First Differences w.r.t. y
Vy_f(1:end-1,:) = (V(2:end,:) - V(1:end-1,:))/dy;
Vy_b(2:end,:) = (V(2:end,:) - V(1:end-1,:))/dy;
Vy_c(2:end-1,:) = 0.5*(Vy_f(2:end-1,:) + Vy_b(2:end-1,:));

% First Differences w.r.t. phi
Vphi_f(:,1:end-1) = (V(:,2:end) - V(:,1:end-1))/dphi;
Vphi_b(:,2:end) = (V(:,2:end) - V(:,1:end-1))/dphi;
Vphi_c(:,2:end-1) = 0.5*(Vphi_f(:,2:end-1) + Vphi_b(:,2:end-1));

%Second Differences w.r.t. y
Vyy = (Vy_f - Vy_b)/dy;

%Vyy(end,:) = 10^(120);

%Second Differences w.r.t. phi
Vphiphi = (Vphi_f - Vphi_b)/dphi;

%Cross Derivative
Vyphi(2:end-1, 2:end-1) = (2*V(2:end-1, 2:end-1) + V(3:end, 3:end) + V(1:end-2, 1:end-2) +...
                              -V(3:end, 2:end-1) - V(1:end-2, 2:end-1) - V(2:end-1, 3:end) - V(2:end-1, 1:end-2))/(2*dphi*dy);
   
% Reduce to original grid
V = V(2:end-1,2:end-1);
Vy_f = Vy_f(2:end-1,2:end-1);
Vy_b = Vy_b(2:end-1,2:end-1);
Vy_c = Vy_c(2:end-1,2:end-1);
Vphi_f = Vphi_f(2:end-1,2:end-1);
Vphi_b = Vphi_b(2:end-1,2:end-1);
Vphi_c = Vphi_c(2:end-1,2:end-1);
Vyy = Vyy(2:end-1,2:end-1); 
Vphiphi = Vphiphi(2:end-1,2:end-1); 
Vyphi = Vyphi(2:end-1,2:end-1); 

%Adjust derivatives at boundary
Vy_b(1,:) = Vy_f(1,:);
Vy_f(end,:) = Vy_b(end,:);

Vy_c(1,:) = Vy_f(1,:);
Vy_c(end,:) = Vy_b(end,:);
            
%% Initialize for optimization
C = zeros(n_grid,n_grid);
Beta = zeros(n_grid,n_grid);
Eta = zeros(n_grid, n_grid);

options = optimoptions('patternsearch',...
            'UseVectorized',true,'Display','off',...
            'ConstraintTolerance', 1e-15,...
            'FunctionTolerance', 1e-15,...
            'MeshTolerance', 1e-15,...
            'MaxIterations', 2000);

%% Find optimizers
parfor iii = 1:n_grid*n_grid
    x = PHI_mat(iii);
    y = Y_mat(iii);
   
    v = V(iii);
    vy_f = Vy_f(iii);
    vy_b = Vy_b(iii);
    vx_c = Vphi_c(iii);
    vy_c = Vy_c(iii);
    vyy = Vyy(iii);
    vxx = Vphiphi(iii);
    vxy = Vyphi(iii);

    c0 = C0(iii);
    beta0 = Beta0(iii);
    eta0 = Eta0(iii);

    Aeq = [];
    Beq = [];

    obj = @(z) optimal_fun(z, x, y, v, vy_f,vy_b, vy_c, vx_c, vyy, vxx, vxy, params);

    zz = patternsearch(obj,[c0, beta0, eta0],[],[],Aeq,Beq,[0,eta_max*(1-x)*vy_f, eta_max],[M,M, eta_max],[],options);

    C(iii) = zz(1);
    Beta(iii) = zz(2);
    Eta(iii) = zz(3);

end


%% Update variables
ExRet = alpha + sigma*Eta.*PHI_mat;

K = (Beta - Eta.*(1-PHI_mat).*Vy_f).*C.^rho./(sigma*lambda);

mu_v = ra - (C.^(1-rho))/(1-rho) + 0.5*rho*Beta.^2;

mu_y = (1-PHI_mat).*ExRet.*Eta.*C.^rho./(sigma*lambda) +...
            Y_mat.*( r-ra + rho/(1-rho).*C.^(1-rho) + 0.5*rho.*Beta.^2)+...
            rho*Eta.*Beta.*PHI_mat.*Y_mat;

f = C - ExRet.*Beta.*C.^rho./(sigma*lambda) + Y_mat.*Eta.*Beta.*PHI_mat;

vol_y = (-rho.*Y_mat.*Beta - Eta.*PHI_mat.*Y_mat);
vol_phi = (PHI_mat.*(1-PHI_mat).*Eta);

end


